
import chisel3._
import chiseltest._
import org.scalatest.{Matchers, FlatSpec}
import scala.util.Random

import comlib.chisel._


class CounterSpec extends FlatSpec with ChiselScalatestTester with Matchers {

    behavior of "Counter"

    val rnd = new Random

    it should "support basic inc counter" in {
        test( new Module {
            val io = IO(new Bundle{
                val cnt = Output(UInt())
            })
            io.cnt := Counter( init = 0.U(4.W) )
        } ){ dut =>
            for( i <- 0 until 100 ){
                val cnt = i % 16
                dut.io.cnt.expect(cnt.U)
                dut.clock.step()
            }
        }
    }

    it should "support combination counter" in {
        test( new Module {
            val io = IO(new Bundle{
                val cnt = Output(UInt())
            })
            val cnt = RegInit(0.U(4.W))
            cnt := Counter(4).load(cnt)
            io.cnt := cnt
        } ){ dut =>
            for( i <- 0 until 100 ){
                val cnt = i % 16
                dut.io.cnt.expect(cnt.U)
                dut.clock.step()
            }
        }
    }

    it should "load initial value" in {
        val init = rnd.nextInt(16)
        test( new Module {
            val io = IO( new Bundle {
                val cnt = Output(UInt(4.W))
            } )

            io.cnt := Counter( init = init.U(4.W) )
        } ){ dut => 
            dut.io.cnt.expect( init.U )
        }
    }

    it should "support inc method" in {
        test( new Module {
            val io = IO(new Bundle{ 
                val inc = Input(Bool())
                val cnt = Output(UInt()) 
            })
            io.cnt := Counter( init=0.U(16.W) ).increase(io.inc)
        } ){ dut =>

            // Inc enable test
            dut.io.inc.poke(true.B)
            val r = rnd.nextInt(20)
            dut.clock.step(r)
            dut.io.cnt.expect( r.U )

            // Inc disable test
            dut.io.inc.poke(false.B)
            dut.clock.step(rnd.nextInt(20))
            dut.io.cnt.expect( r.U )

            // Re-enable
            dut.io.inc.poke(true.B)
            val r1 = rnd.nextInt(20)
            dut.clock.step(r1)
            dut.io.cnt.expect( (r+r1).U )
        }
    }

    it should "support incease method" in {
        test( new Module {
            val io = IO(new Bundle{ 
                val inc = Input(Bool())
                val ival = Input(UInt(16.W))
                val cnt = Output(UInt()) 
            })
            io.cnt := Counter( 16 ).increase(io.inc, io.ival)
        } ){ dut =>

            // Inc enable test
            val step0 = rnd.nextInt(20)
            val ival0 = rnd.nextInt(20)
            dut.io.inc.poke(true.B)
            dut.io.ival.poke( ival0.U )
            dut.clock.step(step0)
            dut.io.cnt.expect( (step0*ival0).U )

            // Inc disable test
            dut.io.inc.poke(false.B)
            dut.clock.step(rnd.nextInt(20))
            dut.io.cnt.expect( (step0*ival0).U )

            // Re-enable
            val step1 = rnd.nextInt(20)
            val ival1 = rnd.nextInt(20)
            dut.io.inc.poke(true.B)
            dut.io.ival.poke( ival1.U )
            dut.clock.step(step1)
            dut.io.cnt.expect( (step0*ival0 + step1*ival1).U )
        }
    }

    it should "support round-robin(loop) count" in {
        test( new Module {
            val io = IO(new Bundle{
                val cnt = Output(UInt())
            })
            io.cnt := Counter(4).increase().value()
        } ){ dut =>
            var r = rnd.nextInt(100)
            var c = r % 16
            dut.clock.step( r )
            dut.io.cnt.expect( c.U )

            // Next
            r = rnd.nextInt(100)
            c = (c + r) % 16
            dut.clock.step(r)
            dut.io.cnt.expect( c.U )
        }
    }

    it should "support increase and decrease at the same time for LOOP2 mode" in {
        test( new Module {
            val io = IO(new Bundle{
                val ival = Input(UInt(4.W))
                val dval = Input(UInt(4.W))
                val cnt = Output(UInt())
            })
            io.cnt := Counter(4).count( true.B, io.ival, true.B, io.dval )
        } ){dut =>
            var (ival, dval, cnt, nxt) = (0, 0, 0, 0)
            for( _ <- 0 until 100 ){
                ival = rnd.nextInt(16)
                dval = rnd.nextInt(16)
                nxt = ( cnt + ival - dval + 16 ) % 16
                dut.io.ival.poke( ival.U )
                dut.io.dval.poke( dval.U )
                dut.clock.step()
                dut.io.cnt.expect( nxt.U, s"ival=${ival}, dval=${dval} and cnt=${cnt}" )
                cnt = nxt
            }
        }
    }

    it should "support clear and so on wrapper" in {
        test( new Module {
            val io = IO(new Bundle{
                val clr = Input(Bool())
                val cnt = Output(UInt())
            })
            io.cnt := Counter( init=0.U(4.W) ).force0(io.clr)
        } ){dut =>
            var cnt = 0
            for( _ <- 0 until 10 ){
                val r = rnd.nextInt(100)
                dut.clock.step(r)
                cnt = (cnt + r) % 16
                dut.io.cnt.expect( cnt.U )
                dut.io.clr.poke(true.B)
                dut.clock.step()
                dut.io.clr.poke(false.B)
                cnt = 0
            }
        }
    }

    it should "support increase and read-clear at same time in count method and inc&clear method" in {
        test( new Module{
            val io = IO( new Bundle{
                val ival = Input(UInt(4.W))
                val clr = Input(Bool())
                val cnt1 = Output(UInt())
                val cnt2 = Output(UInt())
            } )
            io.cnt1 := Counter(4).count( true.B, io.ival, io.clr )
            io.cnt2 := Counter(4).inc(true.B, io.ival).clear(io.clr)
        } ){dut =>
            var cnt = 0
            for( _ <- 0 until 100 ){
                val ival = rnd.nextInt(16)
                val clr = rnd.nextInt(1)==0
                dut.io.ival.poke( ival.U )
                dut.io.clr.poke( clr.B )

                dut.clock.step()

                cnt = if( clr ) ival else ( cnt + ival ) % 16
                dut.io.cnt1.expect( cnt.U )
                dut.io.cnt2.expect( cnt.U )
            }
        }
    }

    it should "support inc method and dec method at the same time" in {
        test( new Module{
            val io = IO( new Bundle{
                val ival = Input(UInt(4.W))
                val dval = Input(UInt(4.W))
                val cnt = Output(UInt())
            } )
            io.cnt := Counter(4).inc(true.B, io.ival).dec(true.B, io.dval)
        } ){dut =>
            dut.io.ival.poke(8.U)
            dut.io.dval.poke(4.U)
            dut.clock.step()
            dut.io.cnt.expect( 4.U )
        }
    }

    // it should "support increase and decrease at the same time for SAT mode" in {
    //     test( new Module {
    //         val io = IO(new Bundle{
    //             val ival = Input(UInt(4.W))
    //             val dval = Input(UInt(4.W))
    //             val cnt = Output(UInt())
    //         })
    //         io.cnt := Counter(2.U(4.W), "SAT").boundary(13.U, 0.U).count( true.B, io.ival, true.B, io.dval ).value()
    //     } ){dut =>
    //         var (ival, dval, cnt, nxt) = (0, 0, 0, 0)
    //         for( _ <- 0 until 100 ){
    //             ival = rnd.nextInt(16)
    //             dval = rnd.nextInt(16)
    //             nxt = (cnt + ival - dval).min(13).max(0)
    //             dut.io.ival.poke( ival.U )
    //             dut.io.dval.poke( dval.U )
    //             dut.clock.step()
    //             dut.io.cnt.expect( nxt.U, s"ival=${ival}, dval=${dval} and cnt=${cnt}" )
    //             cnt = nxt
    //         }
    //     }
    // }

}